# vanna_demo.py
"""
演示如何使用vanna库根据自然语言生成SQL语句。
"""

from vanna.remote import VannaDefault
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
from vanna.flask import VannaFlaskApp
import pandas as pd
from openai import OpenAI
import os
from dotenv import load_dotenv

load_dotenv()
client = OpenAI(
    api_key=os.getenv('DEEPSEEK_API_KEY'),
    base_url=os.getenv('API_BASE_URL', 'https://api.deepseek.com/v1')
)

class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, client=None, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, client=client, config=config)
        self.allow_llm_to_see_sql = True


vn = MyVanna(client=client, config={"model": os.getenv('MODEL_NAME', 'deepseek-chat')})
vn.max_tokens = 1000
vn.temperature = 0.5

vn.connect_to_postgres(
    host=os.getenv('DB_HOST', 'localhost'),
    dbname=os.getenv('DB_NAME', 'postgres'),
    user=os.getenv('DB_USER', 'root'),
    password=os.getenv('DB_PASSWORD', ''),
    port=os.getenv('DB_PORT', 3306)
)

vn.train(
    documentation=open('数据库信息.txt', 'r', encoding='utf-8').read(),
)
# vn.train(question="查询所有租户的基本信息", sql="SELECT * FROM system_tenant")

app = VannaFlaskApp(vn, allow_llm_to_see_data=True)
app.run()

# 示例自然语言问题
# question = "查询所有租户的基本信息"

# # 生成SQL语句
# generated_sql = vn.generate_sql(question)

# print(f"自然语言问题: {question}")
# print(f"生成的SQL语句: {generated_sql}") 